import time
from config import logger, INPUT_DEADZONE
import numpy as np
from collections import deque

class JoystickManager:
    def __init__(self, robot_controller):
        self.joystick_data = []  # 存储历史数据
        self.latest_joystick = None  # 最新数据
        self.robot_controller = robot_controller
    
    def process_joystick_data(self, data):
        """处理接收到的手柄数据"""
        # 添加服务器时间戳
        data['server_time'] = time.time()
        
        # 存储数据
        self.joystick_data.append(data)
        # 只保留最近100条记录
        if len(self.joystick_data) > 100:
            self.joystick_data.pop(0)
            
        # 更新最新数据
        self.latest_joystick = data
        
        # 更新机器人控制器的数据
        self.robot_controller.update_joystick_data(data)
        
        logger.info(f"收到手柄数据: {data}")
    
    def get_latest_data(self):
        """获取最新的手柄数据"""
        return self.latest_joystick
    
    def get_all_data(self):
        """获取所有的手柄数据"""
        return self.joystick_data

class InputSmoother:
    def __init__(self, buffer_size=5, alpha=0.2, deadzone=INPUT_DEADZONE):
        """
        初始化输入平滑处理器
        
        参数:
            buffer_size: 移动平均窗口大小
            alpha: 指数移动平均的权重参数 (0-1之间，越小平滑效果越强)
            deadzone: 死区大小 (小于此值的输入视为0)
        """
        self.buffer_size = buffer_size
        self.alpha = alpha
        self.deadzone = deadzone
        
        # 为各个输入轴初始化缓冲区
        self.buffers = {
            'left_x': deque([0] * buffer_size, maxlen=buffer_size),
            'left_y': deque([0] * buffer_size, maxlen=buffer_size),
            'right_x': deque([0] * buffer_size, maxlen=buffer_size),
            'right_y': deque([0] * buffer_size, maxlen=buffer_size)
        }
        
        # 为各个输入轴初始化EMA值
        self.ema_values = {
            'left_x': 0,
            'left_y': 0,
            'right_x': 0,
            'right_y': 0
        }
        
        # 为各个输入轴初始化当前值(用于惯性模拟)
        self.current_values = {
            'left_x': 0,
            'left_y': 0,
            'right_x': 0,
            'right_y': 0
        }
    
    def apply_deadzone(self, value):
        """应用死区"""
        if abs(value) < self.deadzone:
            return 0
        # 应用死区后重新映射输入范围，保持平滑过渡
        return (value - np.sign(value) * self.deadzone) / (1 - self.deadzone)
    
    def apply_moving_average(self, axis, value):
        """应用移动平均滤波"""
        self.buffers[axis].append(value)
        return sum(self.buffers[axis]) / self.buffer_size
    
    def apply_exponential_moving_average(self, axis, value):
        """应用指数移动平均滤波"""
        self.ema_values[axis] = self.alpha * value + (1 - self.alpha) * self.ema_values[axis]
        return self.ema_values[axis]
    
    def apply_inertia(self, axis, target_value, rate=0.2):
        """
        应用惯性模拟使运动更平滑
        
        参数:
            axis: 输入轴
            target_value: 目标值
            rate: 变化速率 (0-1之间，越小惯性效果越强)
        """
        current = self.current_values[axis]
        # 根据目标值和当前值的差距计算变化量
        delta = (target_value - current) * rate
        new_value = current + delta
        self.current_values[axis] = new_value
        return new_value
    
    def apply_response_curve(self, value, exponent=2.0):
        """
        应用非线性响应曲线
        
        参数:
            value: 输入值 (-1到1之间)
            exponent: 指数值 (>1使小输入更精细，大输入更剧烈)
        """
        return np.sign(value) * (abs(value) ** exponent)
    
    def smooth_input(self, joystick_data, method="combined"):
        """
        平滑处理输入数据
        
        参数:
            joystick_data: 包含手柄输入的字典，包含'left_x', 'left_y', 'right_x', 'right_y'键
            method: 平滑方法，可选值有"moving_avg", "ema", "inertia", "combined"
            
        返回:
            处理后的输入数据字典
        """
        result = {}
        # 确保时间戳等其他字段保留
        for key, value in joystick_data.items():
            if key not in self.buffers:
                result[key] = value
        
        for axis in self.buffers.keys():
            if axis not in joystick_data:
                continue
                
            value = joystick_data[axis]
            
            # 应用死区
            value = self.apply_deadzone(value)
            
            if method == "moving_avg":
                value = self.apply_moving_average(axis, value)
            elif method == "ema":
                value = self.apply_exponential_moving_average(axis, value)
            elif method == "inertia":
                value = self.apply_inertia(axis, value)
            elif method == "combined":
                # 组合多种方法
                value = self.apply_moving_average(axis, value)
                value = self.apply_response_curve(value)
                value = self.apply_inertia(axis, value)
            
            result[axis] = value
            
        return result

# 使用示例
if __name__ == "__main__":
    smoother = InputSmoother()
    
    # 模拟手柄输入序列
    inputs = [
        {'left_x': 0.1, 'left_y': 0.05, 'right_x': 0, 'right_y': 0, 'timestamp': 1},
        {'left_x': 0.2, 'left_y': 0.1, 'right_x': 0, 'right_y': 0, 'timestamp': 2},
        {'left_x': 0.5, 'left_y': 0.3, 'right_x': 0.1, 'right_y': 0.1, 'timestamp': 3},
        {'left_x': 0.7, 'left_y': 0.6, 'right_x': 0.2, 'right_y': 0.2, 'timestamp': 4},
        {'left_x': 0.9, 'left_y': 0.8, 'right_x': 0.3, 'right_y': 0.3, 'timestamp': 5},
    ]
    
    # 处理并打印结果
    for input_data in inputs:
        smoothed = smoother.smooth_input(input_data)
        print(f"原始: {input_data['left_x']:.2f}, {input_data['left_y']:.2f} -> 平滑后: {smoothed['left_x']:.2f}, {smoothed['left_y']:.2f}")
